"""
机器学习预测策略
基于BaseStrategy框架的实现，整合了外部策略文件的优化功能
"""

import numpy as np
import pandas as pd
from typing import Dict, Any, Optional
from datetime import datetime, timedelta
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler

try:
    from ..base_strategy import BaseStrategy, TradingSignal, SignalType
except ImportError:
    # For direct script execution
    import sys
    import os
    sys.path.append(os.path.dirname(os.path.dirname(__file__)))
    from base_strategy import BaseStrategy, TradingSignal, SignalType
from tqsdk.ta import MA, RSI


class MLPredictStrategy(BaseStrategy):
    """
    机器学习预测策略
    
    使用动量、波动性和市场情绪等特征进行机器学习建模，
    通过预测未来价格走势来制定交易决策
    """
    
    def on_initialize(self) -> None:
        """策略初始化"""
        # 策略参数
        self.lookback_window = self.config.get('lookback_window', 100)
        self.prediction_horizon = self.config.get('prediction_horizon', 5)
        self.stop_loss_pct = self.config.get('stop_loss_pct', 0.015)
        self.take_profit_pct = self.config.get('take_profit_pct', 0.03)
        self.retrain_interval = self.config.get('retrain_interval', 20)  # 天
        self.prediction_threshold = self.config.get('prediction_threshold', 0.001)
        
        # 机器学习模型
        self.model = LinearRegression()
        self.scaler = StandardScaler()
        self.is_model_trained = False
        
        # 交易状态
        self.stop_loss_price = 0.0
        self.take_profit_price = 0.0
        self.last_trade_time = None
        self.last_train_time = None
        
        # 数据缓存
        self.kline_data = pd.DataFrame()
        
        self.logger.info(f"机器学习预测策略初始化完成: {self.name}")
    
    def on_bar(self, bar_data: Dict[str, Any]) -> None:
        """K线数据回调"""
        # 更新K线数据
        self._update_kline_data()
        
        # 检查数据是否足够
        if len(self.kline_data) < self.lookback_window + self.prediction_horizon:
            return
        
        # 检查是否需要训练或重新训练模型
        if not self.is_model_trained or self._should_retrain():
            if self._train_model():
                self.last_train_time = datetime.now()
                self.logger.info("模型训练完成")
            else:
                self.logger.warning("模型训练失败")
                return
        
        # 执行交易逻辑
        self._execute_trading_logic()
    
    def on_tick(self, tick_data: Dict[str, Any]) -> None:
        """Tick数据回调"""
        # 实时更新价格用于止损止盈检查
        if self.position != 0:
            current_price = tick_data.get('last_price', 0)
            if current_price > 0:
                self._check_stop_loss_take_profit(current_price)
    
    def on_signal(self) -> Optional[TradingSignal]:
        """信号生成回调"""
        # 信号生成逻辑在_execute_trading_logic中处理
        return None
    
    def on_position_sizing(self, signal: TradingSignal) -> int:
        """仓位计算回调"""
        # 固定仓位大小
        return 1
    
    def _update_kline_data(self) -> None:
        """更新K线数据"""
        if self.data_manager:
            # 获取足够的历史数据用于特征提取和模型训练
            data_length = max(1000, self.lookback_window * 2)
            self.kline_data = self.data_manager.get_kline_data(
                self.symbol, 5*60, data_length  # 5分钟K线
            )
    
    def _should_retrain(self) -> bool:
        """判断是否需要重新训练模型"""
        if self.last_train_time is None:
            return True
        
        days_since_train = (datetime.now() - self.last_train_time).days
        return days_since_train >= self.retrain_interval
    
    def _train_model(self) -> bool:
        """训练机器学习模型"""
        try:
            if len(self.kline_data) < self.lookback_window + self.prediction_horizon:
                return False
            
            # 准备训练数据
            X, y = self._prepare_training_data()
            
            if len(X) == 0 or len(y) == 0:
                return False
            
            # 标准化特征
            X_scaled = self.scaler.fit_transform(X)
            
            # 训练模型
            self.model.fit(X_scaled, y)
            self.is_model_trained = True
            
            self.logger.info(f"模型训练完成，训练样本数: {len(X)}")
            return True
            
        except Exception as e:
            self.logger.error(f"模型训练失败: {e}")
            return False
    
    def _prepare_training_data(self) -> tuple:
        """准备训练数据"""
        features_list = []
        targets_list = []
        
        for i in range(self.lookback_window, len(self.kline_data) - self.prediction_horizon):
            # 获取历史数据片段
            data_slice = self.kline_data.iloc[i-self.lookback_window:i].copy()
            
            # 提取特征
            features = self._extract_features(data_slice)
            if features is not None:
                features_list.append(features)
                
                # 定义目标变量：未来价格变化率
                future_price = self.kline_data['close'].iloc[i + self.prediction_horizon - 1]
                current_price = self.kline_data['close'].iloc[i - 1]
                target = (future_price - current_price) / current_price
                targets_list.append(target)
        
        return np.array(features_list), np.array(targets_list)
    
    def _extract_features(self, kline_data: pd.DataFrame) -> Optional[np.ndarray]:
        """提取特征"""
        try:
            features = []
            
            # 动量特征
            for period in [3, 5, 10, 20, 30]:
                momentum = kline_data['close'].pct_change(period).fillna(0).iloc[-1]
                features.append(momentum)
            
            # 波动性特征
            returns = kline_data['close'].pct_change().fillna(0)
            for period in [5, 10, 20, 30]:
                volatility = returns.rolling(period).std().fillna(0).iloc[-1]
                features.append(volatility)
            
            # RSI特征
            for period in [5, 10, 14, 20]:
                try:
                    rsi_val = RSI(kline_data, period).rsi.iloc[-1]
                    if pd.isna(rsi_val):
                        rsi_val = 50
                    features.append(rsi_val)
                except:
                    features.append(50)
            
            # 均线特征
            for period in [5, 10, 20, 50]:
                try:
                    ma_val = MA(kline_data, period).ma.iloc[-1]
                    if pd.isna(ma_val):
                        ma_val = kline_data['close'].iloc[-1]
                    # 价格相对于均线的位置
                    ma_ratio = kline_data['close'].iloc[-1] / ma_val if ma_val != 0 else 1
                    features.append(ma_ratio)
                except:
                    features.append(1)
            
            # 价格位置特征
            high_20 = kline_data['high'].rolling(20).max().iloc[-1]
            low_20 = kline_data['low'].rolling(20).min().iloc[-1]
            if high_20 != low_20:
                price_position = (kline_data['close'].iloc[-1] - low_20) / (high_20 - low_20)
            else:
                price_position = 0.5
            features.append(price_position)
            
            # 价格变化率特征
            for period in [5, 10]:
                if len(kline_data) > period:
                    price_change = (kline_data['close'].iloc[-1] / kline_data['close'].iloc[-period-1] - 1)
                    features.append(price_change)
                else:
                    features.append(0)
            
            return np.array(features)
            
        except Exception as e:
            self.logger.error(f"特征提取失败: {e}")
            return None
    
    def _predict(self) -> Optional[float]:
        """预测未来价格变化"""
        try:
            if not self.is_model_trained:
                return None
            
            # 提取当前特征
            features = self._extract_features(self.kline_data)
            if features is None:
                return None
            
            # 标准化特征
            features_scaled = self.scaler.transform(features.reshape(1, -1))
            
            # 预测
            prediction = self.model.predict(features_scaled)[0]
            return prediction
            
        except Exception as e:
            self.logger.error(f"预测失败: {e}")
            return None
    
    def _execute_trading_logic(self) -> None:
        """执行交易逻辑"""
        try:
            # 获取预测结果
            prediction = self._predict()
            if prediction is None:
                return
            
            current_price = self.kline_data['close'].iloc[-1]
            current_time = datetime.now()
            
            # 控制交易频率
            if (self.last_trade_time and 
                (current_time - self.last_trade_time).total_seconds() < 300):  # 5分钟间隔
                return
            
            # 计算ATR用于动态止损止盈
            atr_val = self._calculate_atr()
            
            if self.position == 0:  # 无持仓
                if prediction > self.prediction_threshold:
                    # 做多信号
                    signal = TradingSignal(
                        strategy_name=self.name,
                        symbol=self.symbol,
                        signal_type=SignalType.LONG,
                        volume=1,
                        price=current_price,
                        timestamp=current_time,
                        confidence=min(abs(prediction) * 10, 1.0),
                        metadata={'prediction': prediction}
                    )
                    
                    # 设置止损止盈
                    if atr_val > 0:
                        self.stop_loss_price = current_price - 2 * atr_val
                        self.take_profit_price = current_price + 3 * atr_val
                    else:
                        self.stop_loss_price = current_price * (1 - self.stop_loss_pct)
                        self.take_profit_price = current_price * (1 + self.take_profit_pct)
                    
                    # 更新状态
                    self.position = 1
                    self.entry_price = current_price
                    self.last_trade_time = current_time
                    
                    # 触发信号回调
                    if self.signal_callback:
                        self.signal_callback(signal)
                    
                    self.logger.info(f"生成做多信号: 预测={prediction:.4f}, 价格={current_price}")
                
                elif prediction < -self.prediction_threshold:
                    # 做空信号
                    signal = TradingSignal(
                        strategy_name=self.name,
                        symbol=self.symbol,
                        signal_type=SignalType.SHORT,
                        volume=1,
                        price=current_price,
                        timestamp=current_time,
                        confidence=min(abs(prediction) * 10, 1.0),
                        metadata={'prediction': prediction}
                    )
                    
                    # 设置止损止盈
                    if atr_val > 0:
                        self.stop_loss_price = current_price + 2 * atr_val
                        self.take_profit_price = current_price - 3 * atr_val
                    else:
                        self.stop_loss_price = current_price * (1 + self.stop_loss_pct)
                        self.take_profit_price = current_price * (1 - self.take_profit_pct)
                    
                    # 更新状态
                    self.position = -1
                    self.entry_price = current_price
                    self.last_trade_time = current_time
                    
                    # 触发信号回调
                    if self.signal_callback:
                        self.signal_callback(signal)
                    
                    self.logger.info(f"生成做空信号: 预测={prediction:.4f}, 价格={current_price}")
            
            else:  # 有持仓，检查是否需要平仓
                should_close = False
                close_reason = ""
                
                # 检查预测反转
                if ((self.position > 0 and prediction < -self.prediction_threshold) or
                    (self.position < 0 and prediction > self.prediction_threshold)):
                    should_close = True
                    close_reason = "预测反转"
                
                if should_close:
                    self._close_position(current_price, close_reason)
                    
        except Exception as e:
            self.logger.error(f"交易逻辑执行失败: {e}")
    
    def _check_stop_loss_take_profit(self, current_price: float) -> None:
        """检查止损止盈"""
        should_close = False
        close_reason = ""
        
        if self.position > 0:  # 多头持仓
            if current_price <= self.stop_loss_price:
                should_close = True
                close_reason = "多头止损"
            elif current_price >= self.take_profit_price:
                should_close = True
                close_reason = "多头止盈"
        
        elif self.position < 0:  # 空头持仓
            if current_price >= self.stop_loss_price:
                should_close = True
                close_reason = "空头止损"
            elif current_price <= self.take_profit_price:
                should_close = True
                close_reason = "空头止盈"
        
        if should_close:
            self._close_position(current_price, close_reason)
    
    def _close_position(self, current_price: float, reason: str) -> None:
        """平仓"""
        signal = TradingSignal(
            strategy_name=self.name,
            symbol=self.symbol,
            signal_type=SignalType.CLOSE,
            volume=abs(self.position),
            price=current_price,
            timestamp=datetime.now(),
            confidence=1.0,
            metadata={'close_reason': reason}
        )
        
        # 记录交易
        pnl = self._calculate_pnl(current_price)
        self.trades.append({
            'entry_price': self.entry_price,
            'exit_price': current_price,
            'position': self.position,
            'pnl': pnl,
            'close_reason': reason,
            'timestamp': datetime.now()
        })
        
        # 重置状态
        self.position = 0
        self.entry_price = 0.0
        self.stop_loss_price = 0.0
        self.take_profit_price = 0.0
        self.last_trade_time = datetime.now()
        
        # 触发信号回调
        if self.signal_callback:
            self.signal_callback(signal)
        
        self.logger.info(f"{reason}: 价格={current_price}, 盈亏={pnl:.2f}")
    
    def _calculate_pnl(self, exit_price: float) -> float:
        """计算盈亏"""
        if self.position > 0:
            return (exit_price - self.entry_price) * self.position
        elif self.position < 0:
            return (self.entry_price - exit_price) * abs(self.position)
        return 0.0
    
    def _calculate_atr(self, period: int = 14) -> float:
        """计算ATR"""
        try:
            if len(self.kline_data) < period:
                return 0.0
            
            high_low = self.kline_data['high'] - self.kline_data['low']
            high_close = np.abs(self.kline_data['high'] - self.kline_data['close'].shift(1))
            low_close = np.abs(self.kline_data['low'] - self.kline_data['close'].shift(1))
            
            tr = np.maximum(high_low, np.maximum(high_close, low_close))
            atr = tr.rolling(period).mean().iloc[-1]
            
            return atr if not pd.isna(atr) else 0.0
            
        except Exception as e:
            self.logger.error(f"ATR计算失败: {e}")
            return 0.0